from Dialogs import *
from SiteFrequencySpectrum import *
from Parameters import *
from PolymorphismTables import PolymorphismTables
import SiteFrequencySpectrum as sfs
import infiniteSitesTools as ist

from Polyspect_Power import PowerPoint
from Polyspect_Operator import Polyspect
from Polyspect_Calculus import *

import numpy.matlib as np
from scipy.linalg import norm


class NewtonRaphson():
	
	myPath = []
	myParameters = None
	myAltHypothesis = None
	
	myCalculus = None
	mySFS_Null = None
	mySFS_Alt = None
	
	myNumSimulations = 100000
	myNumPowerSimulations = 12500
	
	myTolerance = 0.001
	
	def __init__(self, iStartingPoint, iParameters, iAltHypothesis):
		self.myPath.append( iStartingPoint )
		self.myParameters = iParameters
		self.myAltHypothesis = iAltHypothesis
		
		self.generateSimulationData()
		
		self.myCalculus = PolyspectCalculus( self.myPath[0], \
											self.myParameters, self.myAltHypothesis, \
											self.mySFS_Null, self.mySFS_Alt )
		
		
	def generateSimulationData(self):
		file_Null = './tmp/SFS_Null.ms'
		file_Alt = './tmp/SFS_Alt.ms'
		
		pt = PolymorphismTables()
		msCommand = self.myParameters.getMSString_Null( self.myNumSimulations )
		msOut = pt.runMS(msCommand, file_Null)
		pt.readMS(file_Null, msOut)
	
		self.mySFS_Null = SiteFrequencySpectrum(pt)
		
		pt = PolymorphismTables()
		msCommand = self.myParameters.getMSString_Alternative( self.myAltHypothesis, self.myNumPowerSimulations )
		msOut = pt.runMS(msCommand, file_Alt)
		pt.readMS(file_Alt, msOut)
		
		self.mySFS_Alt = SiteFrequencySpectrum(pt)

	def getMaxPolyspect(self):
		pass
		
	def next(self):
		lastPoint = self.myPath[-1]
		self.myCalculus.setPolyspect( lastPoint )
		
		gradient = self.myCalculus.getDerivative()
		nGradient = norm( gradient )
		print nGradient
		
		if nGradient < self.myTolerance:
			raise StopIteration

		n = self.myParameters.getNumChromosomes()
		
		hessian = self.myCalculus.getConcavity()
		newPoint = np.ravel(np.zeros(n))
		
		for i in range(1, n):
			if abs( hessian[i] ) < 0.01:
				newPoint[i] = lastPoint[i]
			else:
				newPoint[i] = 1.0 * lastPoint[i] - ( gradient[i] / hessian[i] )
			
		self.myPath.append( newPoint )
		return newPoint
	
	def __iter__(self):
		return self

def runTest():
	print "************** NEW RUN ****************"
	print "************** NEW RUN ****************"
	print "************** NEW RUN ****************"

	modelPop = Dialog_MEMES()
	
	memes = modelPop.getResponse()
	
	diaPop = Dialog_Population()
	pop = diaPop.getResponse()

	nullDemography = Demography( pop.N )
	nullDemography.addEpoch( 10000, 0, True )
	print nullDemography
	print pop.getMSString_Null(), nullDemography.getMSString()
	
	demDialog = None
	
	if memes == modelPop.memes[0]:
		demDialog = Dialog_Demography_C( pop.N )
	elif memes == modelPop.memes[1]:
		demDialog = Dialog_Demography_G( pop.N )
	elif memes == modelPop.memes[2]:
		demDialog = Dialog_Demography_CC( pop.N )
	elif memes == modelPop.memes[3]:
		demDialog = Dialog_Demography_CG( pop.N )
	else:
		return
		
	altDemography =  demDialog.getResponse()
		
	simParams = SimulationParameters( pop, nullDemography )
	numChromosomes = simParams.getNumChromosomes()
	linOp = ist.getOp_TajimasD( numChromosomes )
	

	NR = NewtonRaphson( linOp, simParams, altDemography )
	
	for point in NR:
		print "\nPoint is now " + str(point)
		print "\nDerivative is now " + str(NR.myCalculus.getDerivative() )
		print "\nHessian is now " + str(NR.myCalculus.getConcavity() )
		
		

runTest()